import torch
from train_test.utils import collate_fn
from torch.utils.data import DataLoader
from define_dataset import PennFudanDataset
from torchvision.models.detection import fasterrcnn_resnet50_fpn
from data_convert import get_transform
from maskrcnn_netual_rebuild import get_model_instance_segmentation

# model = fasterrcnn_resnet50_fpn(weights="DEFAULT")
# dataset = PennFudanDataset('data/PennFudanPed',get_transform(train=True))
# data_loader = DataLoader(dataset,shuffle=True,batch_size=2,collate_fn=collate_fn)

# #训练
# images, targets = next(iter(data_loader))
# images = list(image for image in images)
# targets = [{k:v for k,v in t.items()} for t in targets]
# output = model(images,targets)

# #返回损失值和预测值
# print(output)

# # For inference
# model.eval()
# x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
# predictions = model(x)  # Returns predictions
# print(predictions[0])


#主函数如下
from train_test.engine import train_one_epoch, evaluate

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

#数据集只有两个类 -- 背景和人
num_classes = 3

#用我们的数据集以及自定义的转换函数
dataset = PennFudanDataset('data/PennFudanPed',get_transform(train=True))
dataset_test = PennFudanDataset('data/PennFudanPed',get_transform(train=False))

#将数据集分成训练数据集和测试数据级
indics = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset,indics[:-50])
dataset_test = torch.utils.data.Subset(dataset_test,indics[-50:])

#定义dataloader
train_dataloader = DataLoader(dataset, batch_size=2, shuffle=True, collate_fn=collate_fn)
test_dataloader = DataLoader(dataset_test, batch_size=1, shuffle=False, collate_fn=collate_fn)

model = get_model_instance_segmentation(num_classes)
model.to(device)

#构建损失函数和优化器
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(
    params,
    lr=0.005,
    momentum=0.9,
    weight_decay=0.0005
)

#添加一个学习率的调度器
lr_scheduler = torch.optim.lr_scheduler.StepLR(
    optimizer,
    step_size=3,
    gamma=0.1
)

num_epochs = 2
for epoch in range(num_epochs):
    train_one_epoch(model,optimizer,train_dataloader,device,epoch,print_freq=10)
    lr_scheduler.step()
    evaluate(model,test_dataloader,device=device)

torch.save(model,"trainmodel.pth")
print("That's it")